# -*- coding: UTF-8 -*-
import json
import time
import sys
import traceback
import requests
import socket
from datetime import datetime
from time import sleep
from threading import Lock, Thread
import websocket
import zlib
import redis
import pymongo
import threading
import queue
from okex.spot_api import SpotAPI
from okex.futures_api import FutureAPI
import schedule
from okex.model.position import PositionModel
from okex.utils import *
import gc
#client = pymongo.MongoClient(host='127.0.0.1', port=27017,username="admin",password="123456")
#db = client['okex_trade']
#collection = db['futures_trade']

# pool = redis.ConnectionPool(host='127.0.0.1', port='6379', db=8,password="123456")
# r = redis.Redis(connection_pool=pool)
# pipe = r.pipeline()
GMT_FORMAT = '%a, %d %b %Y %H:%M:%S GMT+0800 (CST)'
class OKEXWebsocket(object):
    """
        Websocket API
        创建Websocket client对象后，需要调用Start的方法去启动workder和ping线程
        1. Worker线程会自动重连.
        2. 使用stop方法去停止断开和销毁websocket client,
        3. 四个回调方法..
        * on_open
        * on_close
        * on_msg
        * on_error

        start()方法调用后，ping线程每隔60秒回自动调用一次。

    """

    def __init__(self, host=None, ping_interval=5,webclients=[],
    mongo_client=None,redis_client=None,msg_callback=None,
    api_key='b0b114f3-96aa-4f03-9f83-1207a8e771a3',seceret_key='EDEF91B4A52E9B6C821786B022112669',passphrase='TG159357',
    logger=None):
        """Constructor"""
        self.logger = logger
        self.db = mongo_client['okex_trade']
        # 期货交易表
        self.table_future_trade = self.db['futures_trade']
        # 现货交易表
        self.table_spots_trade = self.db['spots_trade']
        # 现货深度表
        self.table_spots_depth = self.db['spots_depth']
        # 期货深度表
        self.table_futures_depth = self.db['futures_depth']
         # 持仓张数表
        self.table_position = self.db['position']
        # 永续
        self.table_swap_depth = self.db['swap_depth']
        self.table_swap_trade = self.db['swap_trade']
        self.redis_client = redis_client
        # 现货API
        self.Okex_spotAPI = SpotAPI(api_key, seceret_key, passphrase, True)
        # 期货API
        self.Okex_futureAPI = FutureAPI(api_key, seceret_key, passphrase, True)

        self.msg_callback = msg_callback

        self.host = host
        self.ping_interval = ping_interval

        self._ws_lock = Lock()
        self._ws = None

        self._worker_thread = None
        self._ping_thread = None
        self._save_db_thread = None
        self._position_thread = None
        self._instrument_thread = None
        self._active = False  # 开启启动websocket的开关。

        # debug需要..
        self._last_sent_text = None
        self._last_received_text = None

        self._webclients = webclients
        self.mongo_client = mongo_client
        # self.fdepth_save_list = queue.PriorityQueue()
        # self.depth_save_list = queue.PriorityQueue()
        self.hold_save_list = queue.PriorityQueue()
        self.instruments = []
        self.data_list = []
        self.fdepth = []
        self.depth = []
        self.swap_depth = []
        self.ftrade = []
        self.trade = []
        # [{contractType: "spot", price: "199.46", updatetime: "Thu, 21 May 2020 15:51:45 GMT+0800 (CST)"}]
        self.topdata = {};

    def get_instruments(self):
        while self._active:
            self.logger.info("更新合约列表")
            url_instruments = 'https://www.okex.com/api/futures/v3/instruments'
            req = requests.get(url_instruments, timeout=self.ping_interval)
            d = json.loads(req.text)
            if d != self.instruments: #合约有变化就重新订阅数据
                self.instruments = json.loads(req.text)
                #订阅数据
                self.on_open()
            time.sleep(3600*4)


    def start(self):
        """
        启动客户端，客户端连接成功后，会调用 on_open这个方法
        on_open 方法调用后，才可以向服务器发送消息的方法.
        """
        self.logger.info("start subscribe okex websockets")
        self._active = True
        self._worker_thread = Thread(target=self._run)
        self._worker_thread.start()

        #保存持仓信息
        self._position_thread = Thread(target=self.get_position)
        self._position_thread.start()

        #更新持仓列表
        self._instrument_thread = Thread(target=self.get_instruments)
        self._instrument_thread.start()

        self._ping_thread = Thread(target=self._run_ping)
        self._ping_thread.start()

        # self._save_db_thread = Thread(target=self._save_hold)
        # self._save_db_thread.start()

        # self._save_depth_thread = Thread(target=self._save_depth)
        # self._save_depth_thread.start()

        # self._save_fdepth_thread = Thread(target=self._save_fdepth)
        # self._save_fdepth_thread.start()

        self.db_worker()

    def stop(self):
        """
        停止客户端.
        """
        self._active = False
        self._disconnect()

    def join(self):
        """
        Wait till all threads finish.
        This function cannot be called from worker thread or callback function.
        """
        self._ping_thread.join()
        self._worker_thread.join()
        # self._save_db_thread.join()
        self._position_thread.join()
        self._instrument_thread.join()

    def get_position(self):
        while self._active:
            for instrument in self.instruments:
                data = self.Okex_futureAPI.get_holds(instrument['instrument_id'])
                d = {"type":"position","data":data}
                j = json.dumps(d)
                self.hold_save_list.put(j)
      
            for symbol in self.topdata:
                if len(self.topdata[symbol]) >= 6:
                    # self.logger.info(self.topdata[symbol])
                    self.redis_client.set("topdata:{0}".format(symbol),json.dumps(self.topdata[symbol]))
            time.sleep(5)

    def send_msg(self, msg: dict):
        """
        向服务器发送数据.
        如果你想发送非json数据，可以重写该方法.
        """
        text = json.dumps(msg)
        #print("send msg:",text)
        self._record_last_sent_text(text)
        return self._send_text(text)

    def _send_text(self, text: str):
        """
        发送文本数据到服务器.
        """
        ws = self._ws
        if ws:
            ws.send(text, opcode=websocket.ABNF.OPCODE_TEXT)

    def _ensure_connection(self):
        """"""
        triggered = False
        with self._ws_lock:
            if self._ws is None:
                self._ws = websocket.create_connection(self.host)

                # triggered = True
        if triggered:
            self.on_open()

    def _disconnect(self):
        """
        """
        triggered = False
        with self._ws_lock:
            if self._ws:
                ws: websocket.WebSocket = self._ws
                self._ws = None

                triggered = True
        if triggered:
            ws.close()
            self.on_close()

    def _run(self):
        """
        保持运行，知道stop方法调用.
        """
        try:
            while self._active:
                try:
                    self._ensure_connection()
                    ws = self._ws
                    if ws:
                        text = ws.recv()
                        # ws object is closed when recv function is blocking
                        if not text:
                            self._disconnect()
                            continue

                        self._record_last_received_text(text)

                        self.on_msg(text)
                # ws is closed before recv function is called
                # For socket.error, see Issue #1608
                except (websocket.WebSocketConnectionClosedException, socket.error):
                    self._disconnect()

                # other internal exception raised in on_msg
                except:  # noqa
                    et, ev, tb = sys.exc_info()
                    self.on_error(et, ev, tb)
                    self._disconnect()  #

        except:  # noqa
            et, ev, tb = sys.exc_info()
            self.on_error(et, ev, tb)

        self._disconnect()

    def _run_ping(self):
        """"""
        while self._active:
            try:
                self._ping()
            except:  # noqa
                et, ev, tb = sys.exc_info()
                self.on_error(et, ev, tb)
                sleep(1)

            for i in range(self.ping_interval):
                if not self._active:
                    break
                sleep(i)

    def db_worker(self):
        workers = [threading.Thread(target=self._save_hold)]
        for w in workers:
            w.setDaemon(True)
            w.start()
        self.hold_save_list.join()
        # self.depth_save_list.join()
        # self.fdepth_save_list.join()
    def _get_instrument_alias(self,instrument_id):
        for i in self.instruments:
            if i['instrument_id'] == instrument_id:
                return i['alias']
        return ''
    def _save_hold(self):
        """"""
        while self._active:
            try:
                jb = self.hold_save_list.get()
                next_job = json.loads(jb)
                data = next_job['data']
                if next_job['type'] == "position": # 持仓
                    data['time'] = get_time(data["timestamp"])
                    arr = data["instrument_id"].split('-')
                    data['symbol'] = arr[0]
                    self.table_position.insert_one(data)
                    self.logger.debug("保存持仓数据:{0}".format(data))
                self.hold_save_list.task_done()
                time.sleep(0.1)
            except:  # noqa
                et, ev, tb = sys.exc_info()
                self.logger.debug("{0}{1}{2}".format(et,ev,tb))
    def _save_fdepth(self):
        """"""
        while self._active:
            try:
                jb = self.fdepth_save_list.get()
                next_job = json.loads(jb)
                data = next_job['data']
                self.logger.info("接收数据:{0}".format(data['instrument_id']))
                if next_job['type'] == "fdepth": # 期货深度
                    data['time'] = get_time(data["timestamp"])
                    arr = data["instrument_id"].split('-')
                    data['symbol'] = arr[0]
                    data['contractType'] = self._get_instrument_alias(data["instrument_id"])
                    if data['contractType'] == '':
                        continue
                    self.table_futures_depth.insert_one(data)
                    self.logger.debug("保存期货深度数据:{0}".format(data))
                self.fdepth_save_list.task_done()
                time.sleep(0.1)
            except:  # noqa
                et, ev, tb = sys.exc_info()
                self.logger.debug("{0}{1}{2}".format(et,ev,tb))
    def _save_depth(self):#保存现货 深度队列
        """"""
        while self._active:
            try:
                jb = self.depth_save_list.get()
                next_job = json.loads(jb)
                data = next_job['data']
                if next_job['type'] == "depth": # 现货深度
                    data['time'] = get_time(data["timestamp"])
                    arr = data["instrument_id"].split('-')
                    data['symbol'] = arr[0]
                    data['contractType'] = 'spot'
                    self.table_spots_depth.insert_one(data)
                    self.logger.debug("保存现货深度数据:{0}".format(data))
                time.sleep(0.1)
            except:  # noqa
                self.depth_save_list.task_done()
                et, ev, tb = sys.exc_info()
                self.logger.info("{0}{1}{2}".format(et,ev,tb))

    def _ping(self):
        """"""
        ws = self._ws
        if ws:
            ws.send("ping", websocket.ABNF.OPCODE_PING)

    def on_open(self):
        time.sleep(5)
        spot_trades = []
        okex_instruments = []
        klines = []
        depths = []
        fdepth = []
        swap_depth = []
        for i in self.instruments:
            if i['quote_currency'] == 'USD': #USD 币本位 USDT 钱本位
                instrument_id = 'futures/trade:' + i['instrument_id']
                #订阅期货
                okex_instruments.append(instrument_id)
                # 同步订阅 现货 去掉后面日期 BTC-USD-200515 => BTC-USDT 现货对应是USDT
                arr = i['instrument_id'].split("-")
                spot_instrument_id = "spot/trade:{0}-{1}T".format(arr[0],arr[1])
                if spot_instrument_id not in spot_trades:
                    spot_trades.append(spot_instrument_id)

                # 订阅k线
                kline_sub = "swap/candle60s:{0}-SWAP".format(spot_instrument_id)
                if kline_sub not in klines:
                    klines.append(kline_sub)
                #订阅 期货深度
                fd = "futures/depth5:"+i['instrument_id']
                fdepth.append(fd)
                # 订阅现货深度
                shendu_sub = "spot/depth5:{0}-{1}T".format(arr[0],arr[1])
                if shendu_sub not in depths:
                    depths.append(shendu_sub)
                # 订阅永续深度
                swap_shendu = "swap/depth5:{0}-{1}-SWAP".format(arr[0],arr[1])
                if swap_shendu not in swap_depth:
                    swap_depth.append(swap_shendu)


        """on open """
        # self.logger.info("subscribe {0}".format(okex_instruments))
        # 订阅交易数据
        data = {"op": "subscribe", "args":okex_instruments}
        self.send_msg(data)
        # 订阅现货交易数据
        data = {"op": "subscribe", "args":spot_trades}
        self.send_msg(data)

        #订阅kline
        # self.logger.info("subscribe candle60s")
        # data = {"op": "subscribe", "args":klines}
        # self.send_msg(data)

        self.logger.info("subscribe depth")
        data = {"op": "subscribe", "args":depths}
        self.send_msg(data)

        self.logger.info("subscribe fdepth")
        data = {"op": "subscribe", "args":fdepth}
        self.send_msg(data)

        self.logger.info("subscribe swap depth")
        data = {"op": "subscribe", "args":swap_depth}
        self.send_msg(data)


    def on_close(self):
        """
        on close websocket
        """

    def on_msg(self, data: str):
        """call when the msg arrive."""
        decompress = zlib.decompressobj(
            -zlib.MAX_WBITS  # see above
        )
        sendStruct = None
        msg = json.loads(decompress.decompress(data))
        if 'table' in msg:
            # self.logger.info(msg['table'])
            pass

        if 'table' in msg and msg['table']=='swap/candle60s': #kline 数据
            data = msg['data'][0]
            sendStruct = {"type":"kline","data":data}

        if 'table' in msg and msg['table']=='spot/depth5': # 现货深度
            data = msg['data'][0]
            sendStruct = {"type":"depth","data":data}
            #现货深度数据
            # self.depth_save_list.put(json.dumps(sendStruct))
            # self.logger.info("条现货深度数据{0}".format(len(self.depth)))
            if len(self.depth) >= 1000:
                self.logger.info("1000 条现货深度数据")
                self.table_spots_depth.insert_many(self.depth)
                self.depth = []
            else:
                data['time'] = get_time(data["timestamp"])
                arr = data["instrument_id"].split('-')
                data['symbol'] = arr[0]
                # self.logger.info(data)
                index = -1 
                if arr[0] not in self.topdata:
                    self.topdata[arr[0]] = []
                for i in range(len(self.topdata[arr[0]])):
                    if self.topdata[arr[0]][i]['contractType'] == "spot":
                        index = i
                        break
                t = data['time']
                localtime=time.localtime(t)
                d = {
                    'contractType': "spot",
                    "price": data['bids'][0][0],
                    "updatetime": time.strftime(GMT_FORMAT, localtime)
                }
                if index > -1:
                    self.topdata[arr[0]][index] = d
                else:
                    self.topdata[arr[0]].append(d)
                self.depth.append(data)

        if 'table' in msg and msg['table']=='swap/depth5': # 永续深度
            data = msg['data'][0]
            sendStruct = {"type":"depth","data":data}
            #现货深度数据
            # self.depth_save_list.put(json.dumps(sendStruct))
            # self.logger.info("条现货深度数据{0}".format(len(self.depth)))
            if len(self.swap_depth) >= 1000:
                self.logger.info("1000 条永续深度数据")
                self.table_swap_depth.insert_many(self.swap_depth)
                self.swap_depth = []
            else:
                data['time'] = get_time(data["timestamp"])
                arr = data["instrument_id"].split('-')
                data['symbol'] = arr[0]
                data['contractType'] = "swap"
                if arr[0] not in self.topdata:
                    self.topdata[arr[0]] = []
                index = -1 
                for i in range(len(self.topdata[arr[0]])):
                    if self.topdata[arr[0]][i]['contractType'] == "swap":
                        index = i
                        break
                t = data['time']
                localtime=time.localtime(t)
                d = {
                    'contractType': "swap",
                    "price": data['asks'][0][0],
                    "updatetime": time.strftime(GMT_FORMAT, localtime)
                }
                if index > -1:
                    self.topdata[arr[0]][index] = d
                else:
                    self.topdata[arr[0]].append(d)
                # self.logger.info(data)
                self.swap_depth.append(data)

        if 'table' in msg and msg['table']=='futures/depth5': # 期货深度
            data = msg['data'][0]
            sendStruct = {"type":"fdepth","data":data}
            #期货深度数据
            # self.logger.info("推送数据 {0}".format(sendStruct['data']['instrument_id']))
            # self.fdepth_save_list.put(json.dumps(sendStruct))
            # self.logger.info("条期货深度数据{0}".format(len(self.fdepth)))
            if len(self.fdepth) >= 2000:
                self.logger.info("2000 条期货深度数据")
                self.table_futures_depth.insert_many(self.fdepth)
                self.fdepth = []
            else:
                data['time'] = get_time(data["timestamp"])
                arr = data["instrument_id"].split('-')
                data['symbol'] = arr[0]
                # self.logger.info(data)
                data['contractType'] = self._get_instrument_alias(data["instrument_id"])
                if data['contractType'] == '':
                    return
                if arr[0] not in self.topdata:
                    self.topdata[arr[0]] = []
                index = -1 
                for i in range(len(self.topdata[arr[0]])):
                    if self.topdata[arr[0]][i]['contractType'] == data['contractType']:
                        index = i
                        break
                t = data['time']
                localtime=time.localtime(t)
                d = {
                    'contractType': data['contractType'],
                    "price": data['bids'][0][0],
                    "updatetime": time.strftime(GMT_FORMAT, localtime)
                }
                if index > -1:
                    self.topdata[arr[0]][index] = d
                else:
                    self.topdata[arr[0]].append(d)
                # self.logger.info(data)
                self.fdepth.append(data)

        if 'table' in msg and msg['table'] == 'futures/trade':#期货交易记录
            data = msg['data'][0]
            sendStruct = {"type":"trade","data":data}
            # self.logger.info("条期货交易数据{0}".format(len(self.ftrade)))
            if len(self.ftrade) >= 1000:
                self.logger.info("1000 条期货交易数据")
                self.table_future_trade.insert_many(self.ftrade)
                self.ftrade = []
            else:
                data['time'] = get_time(data["timestamp"])
                arr = data["instrument_id"].split('-')
                data['symbol'] = arr[0]
                data['contractType'] = self._get_instrument_alias(data["instrument_id"])
                if data['contractType'] == '':
                    return
                self.ftrade.append(data)

        if 'table' in msg and msg['table'] == 'spot/trade':#现货交易记录
            data = msg['data'][0]
            sendStruct = {"type":"trade","data":data}
            # self.logger.info("条现货交易数据{0}".format(len(self.trade)))
            if len(self.trade) >= 1000:
                self.logger.info("1000 条现货交易数据")
                self.table_spots_trade.insert_many(self.trade)
                self.trade = []
            else:
                data['time'] = get_time(data["timestamp"])
                arr = data["instrument_id"].split('-')
                data['symbol'] = arr[0]
                self.trade.append(data)

        if sendStruct != None: # 消息给推送给页面
            self.msg_callback(sendStruct)

    def on_error(self, exception_type: type, exception_value: Exception, tb):
        """
        Callback when exception raised.
        """
        sys.stderr.write(
            self.exception_detail(exception_type, exception_value, tb)
        )

        return sys.excepthook(exception_type, exception_value, tb)

    def exception_detail(
            self, exception_type: type, exception_value: Exception, tb
    ):
        """
        Print detailed exception information.
        """
        text = "[{}]: Unhandled WebSocket Error:{}\n".format(
            datetime.now().isoformat(), exception_type
        )
        text += "LastSentText:\n{}\n".format(self._last_sent_text)
        text += "LastReceivedText:\n{}\n".format(self._last_received_text)
        text += "Exception trace: \n"
        text += "".join(
            traceback.format_exception(exception_type, exception_value, tb)
        )
        return text

    def _record_last_sent_text(self, text: str):
        """
        Record last sent text for debug purpose.
        """
        self._last_sent_text = text[:1000]

    def _record_last_received_text(self, text: str):
        """
        Record last received text for debug purpose.
        """
        self._last_received_text = text[:1000]
